#include <vector>

#include "caffe/layers/euclidean_loss_layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

    template<typename Dtype>
    void EuclideanLossLayer<Dtype>::Reshape(
            const vector<Blob<Dtype> *> &bottom, const vector<Blob<Dtype> *> &top) {
        LossLayer<Dtype>::Reshape(bottom, top);
        CHECK_EQ(bottom[0]->count(1), bottom[1]->count(1))
            << "Inputs must have the same dimension.";
        diff_.ReshapeLike(*bottom[0]);
    }

    template<typename Dtype>
    void EuclideanLossLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype> *> &bottom,
                                                const vector<Blob<Dtype> *> &top) {
        int count = bottom[0]->count();
        caffe_sub(
                count,
                bottom[0]->cpu_data(),
                bottom[1]->cpu_data(),
                diff_.mutable_cpu_data());
//        const Dtype *label_ptr=bottom[1]->cpu_data();
//        const Dtype *data_ptr=bottom[0]->cpu_data();
//        for (int i = 0; i < count; ++i)
//        {
//            if (label_ptr[i] > 0.0000f)
//            {
//                std::cout << "label_data[i]:" << label_ptr[i] << " data_ptr[i]" << data_ptr[i] << std::endl;
//            }
//        }
        Dtype dot = caffe_cpu_dot(count, diff_.cpu_data(), diff_.cpu_data());
        Dtype loss = dot / bottom[0]->num() / Dtype(2);
        top[0]->mutable_cpu_data()[0] = loss;
    }

    template<typename Dtype>
    void EuclideanLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype> *> &top,
                                                 const vector<bool> &propagate_down,
                                                 const vector<Blob<Dtype> *> &bottom) {
        for (int i = 0; i < 2; ++i)
        {
            if (propagate_down[i])
            {
                const Dtype sign = (i == 0) ? 1 : -1;
                const Dtype alpha = sign * top[0]->cpu_diff()[0] / bottom[i]->num();
                caffe_cpu_axpby(
                        bottom[i]->count(),              // count
                        alpha,                              // alpha
                        diff_.cpu_data(),                   // a
                        Dtype(0),                           // beta
                        bottom[i]->mutable_cpu_diff());  // b
            }
        }
    }

#ifdef CPU_ONLY
    STUB_GPU(EuclideanLossLayer);
#endif

    INSTANTIATE_CLASS(EuclideanLossLayer);

    REGISTER_LAYER_CLASS(EuclideanLoss);

}  // namespace caffe
